#include <kern/vm.h>
#include <kern/proc.h>
#include <kern/debug.h>
#include <kern/mmu.h>
#include <types.h>
#include <error.h>
#include <x86.h>
#include <string.h>
#include <elf.h>

/* global vars */
struct pool_t free_procs;

extern void trapret(void) __attribute__ ((noreturn));

// TODO move this function to mp.c and implement it
int cpunum() {
	return 0;
}

static void forkret(void)
{
	monitor_write("about to run for the first time\n");

}

static void init_kstack(struct proc_t* p)
{
	uint8_t* sp;
	extern char bootstacktop[];

	ASSERT(p->kstack != NULL, "kstack is null");

	sp = (uint8_t*)p->kstack + KSTKSIZE;

	sp -= sizeof *p->tf;
	p->tf = (struct trapframe_t*)sp;

	// return to trapret from forkret
	sp -= sizeof (uintptr_t);
	*(uintptr_t*)sp = (uintptr_t)trapret;

	// start at forkret
	sp -= sizeof *p->kctx;
	p->kctx = (struct context_t*)sp;
	memset(p->kctx, 0, sizeof *p->kctx);
	
	p->kctx->eip = (uintptr_t)forkret;
	monitor_write("eip = ");
	monitor_write_hex(p->kctx->eip);
	monitor_write("\n");
	// setup segment registers for user space
	p->tf->ds = (SEG_UDATA << 3) | DPL_USER;
	p->tf->es = (SEG_UDATA << 3) | DPL_USER;
	p->tf->ss = (SEG_UDATA << 3) | DPL_USER;
	p->tf->cs = (SEG_UCODE << 3) | DPL_USER;
	p->tf->esp = USTACKTOP;

	// enable interrupts in user mode
	// TODO: handle interrupts and uncomment this
	p->tf->eflags = FL_IF;

	// eip of tf will be set later
}

// to be used for loading user programs during system startup
static void alloc_segment(pde_t* pgdir, uintptr_t start_va, size_t sz)
{
	uintptr_t va;

	monitor_write("alloc seg: ");
	monitor_write_hex(start_va);
	monitor_write("  ");
	monitor_write_dec(sz);
	monitor_write("\n");

	sz = ROUNDUP(start_va + sz, PGSIZE) - ROUNDDOWN(start_va, PGSIZE);

	//sz = ROUNDUP(start_va - ROUNDDOWN(start_va, PGSIZE) + sz, PGSIZE);
	
	for (va = ROUNDDOWN(start_va, PGSIZE); sz; sz -= PGSIZE, va += PGSIZE) {
		if (vm_is_mapped(pgdir, va)) {
			monitor_write("va ");
			monitor_write_hex(va);
			monitor_write(" is already mapped to ");
			monitor_write_hex(vm_va2pa(pgdir, va));
			monitor_write("\n");
			continue;
		}
		if (vm_insert_page(pgdir, va, PTE_U | PTE_W))
			panic("unable to allocate segment");
		monitor_write("pa = ");
		monitor_write_hex(vm_va2pa(pgdir, va));
		monitor_write("\n");
	}
}

// to be used for loading user programs during system startup
void proc_load_code(struct proc_t* proc, uint8_t* binary, int sz)
{
	struct Elf *elf = (struct Elf *)binary;
	struct Proghdr *ph, *eph;
	
	if (elf->e_magic != ELF_MAGIC)
		panic("Invalid elf format");

	
	ph = (struct Proghdr *)(binary + elf->e_phoff);
	eph = ph + elf->e_phnum;
	for (;ph < eph; ++ph) {
		if (ph->p_type != ELF_PROG_LOAD) 
			continue;
		alloc_segment(proc->pgdir, (uintptr_t)ph->p_va, ph->p_memsz);
		
		ASSERT(ph->p_filesz <= ph->p_memsz, "invalid elf filesz");

		memset((void*)ph->p_va, 0, ph->p_memsz);
		memmove((void*)ph->p_va, binary + ph->p_offset, ph->p_filesz);
	}
	// setup the entry point
	proc->tf->eip = elf->e_entry;
	//proc->tf->eip = (uint32_t)imhere;

	monitor_write("start at: ");
	monitor_write_hex(proc->tf->eip);
	monitor_write("\n");
}

void proc_init(void)
{
	int i;
	pool_init(&free_procs, offsetof(struct proc_t, link));
	for (i = 0; i < NRPROCS; ++i) {
		procs[i].state = PROC_FREE;
		procs[i].id = 0xF000 | i;
		pool_add(&free_procs, &procs[i]);
	}
}

int proc_alloc(struct proc_t** p_proc, pid_t parent)
{
	struct proc_t* proc;
	physaddr_t pa;
	int rc;
	pid_t pid;

	if (!(proc = (struct proc_t*)pool_pop(&free_procs)))
		return -ENOPROC;

	ASSERT(proc->state == PROC_FREE, "non free proc in free pool");

	// pid is the only member that does not change
	pid = proc->id;
	memset(proc, 0, sizeof *proc);
	proc->id = pid;

	if (!(proc->kstack = vm_kalloc_page(&pa))) {
		rc = -ENOMEM;
		goto err;
	}

	if (!(proc->pgdir = vm_kalloc_page(&pa))) {
		rc = -ENOMEM;
		goto err;
	}
		
	vm_init_pgdir(proc->pgdir);
	init_kstack(proc);

	// allocate a user stack
	if ((rc = vm_insert_page(proc->pgdir, USTACKTOP-PGSIZE, PTE_U | PTE_W)))
		goto err;
	
	proc->state = PROC_EMBRYO;
	proc->parent_id = parent;

	*p_proc = proc;

	return 0;
err:
	if (proc->kstack)
		vm_release_kvpage(proc->kstack);
	if (proc->pgdir) {
		if (vm_is_mapped(proc->pgdir, USTACKTOP))
			vm_unmap(proc->pgdir, USTACKTOP);
		vm_release_kvpage(proc->pgdir);
	}
	return rc;
}

int proc_free(struct proc_t* proc)
{
	// TODO free user virtual mem
	proc->state = PROC_FREE;
	pool_add(&free_procs, proc);
	
	return 0;
}






